package com.gitee.wsl.ext.array.sort

import com.gitee.wsl.ext.array.swap
import kotlin.math.min


private fun swap(array: IntArray, index1: Int, index2: Int) {
    val temp = array[index1]
    array[index1] = array[index2]
    array[index2] = temp
}

fun quickSort(array: IntArray, low: Int, high: Int) {
    if (low < high) {
        val j = partition(array, low, high)
        quickSort(array, low, j)
        quickSort(array, j + 1, high)
    }
}

private fun partition(array: IntArray, low: Int, high: Int): Int {
    val pivot = array[low]
    var i = low
    var j = high

    // stop when ----ji----
    while (i < j) {
        // keep incrementing i till you find the element larger than pivot
        while (pivot >= array[i]) {
            i++
        }
        // keep incrementing j till you find the element smaller than pivot
        while (pivot < array[j]) {
            j--
        }

        // swap element at i and j
        if (i < j) {
            swap(array, i, j)
        }
    }
    // swap pivot and j
    swap(array, low, j)
    // return partition index
    return j
}

internal fun<T> Array<T>.med3(idxA: Int, idxB: Int, idxC: Int, comparator: Comparator<T>): Int {
    val ab = comparator.compare(this[idxA], this[idxB])
    val ac = comparator.compare(this[idxA], this[idxC])
    val bc = comparator.compare(this[idxB], this[idxC])

    return when {
        ab < 0 -> {
            when {
                bc < 0 -> idxB
                ac < 0 -> idxC
                else -> idxA
            }
        }

        bc > 0 -> idxB
        ac > 0 -> idxC
        else -> idxA
    }
}

private const val SMALL = 7
private const val MEDIUM = 50

private fun<T> Array<T>.vecSwap(idxA: Int, idxB: Int, n: Int) {
    var a = idxA
    var b = idxB
    repeat(n) {
        this.swap(a++, b++)
    }
}

fun<T> Array<T>.quickSort(fromIdx: Int, toIdx: Int, comparator: Comparator<T>) {
    val len = toIdx - fromIdx

    // Selection sort on smallest arrays
    if (len < SMALL) {
        this.selectionSort(fromIdx, toIdx, comparator)
        return
    }

    // Choose a partition element, v
    var m = fromIdx + len / 2 // Small arrays, middle element
    if (len > SMALL) {
        var l = fromIdx
        var n = toIdx - 1
        if (len > MEDIUM) {
            // Big arrays, pseudo median of 9
            val s = len / 8
            l = this.med3(l, l + s, l + 2 * s, comparator)
            m = this.med3(m - s, m, m + s, comparator)
            n = this.med3(n - 2 * s, n - s, n, comparator)
        }
        // Mid-size, med of 3
        m = this.med3(l, m, n, comparator)
    }

    val v = this[m]
    // Establish Invariant: v* (<v)* (>v)* v*
    var a = fromIdx
    var b = a
    var c = toIdx - 1
    var d = c
    while (true) {
        var comparison = 0

        while (b <= c && comparator.compare(this[b], v).also { comparison = it } <= 0) {
            if (comparison == 0) {
                this.swap(a++, b)
            }
            b++
        }

        while (c >= b && comparator.compare(this[c], v).also { comparison = it } >= 0) {
            if (comparison == 0) {
                this.swap(c, d--)
            }
            c--
        }

        if (b > c) {
            break
        }

        this.swap(b++, c--)
    }

    // Swap partition elements back to the middle
    var s = min(a - fromIdx, b - a)
    this.vecSwap(fromIdx, b - s, s)
    s = min(d - c, toIdx - d - 1)
    this.vecSwap(b, toIdx - s, s)
    // Recursively sort non-partition-elements
    if ((b - a).also { s = it } > 1) {
        this.quickSort(fromIdx, fromIdx + s, comparator)
    }
    if ((d - c).also { s = it } > 1) {
        this.quickSort(toIdx - s, toIdx, comparator)
    }
}
